#include "stdafx.h"
#include "ylLog.h"
#include "ylStringUtils.h"

#include <stdio.h>
#include <iostream>

#include <sys/types.h>
#include <WinSock2.h>
#include <WS2tcpip.h>

#include <DbgHelp.h>

using namespace yoctolog;


Log *Log::ms_pInstance = NULL;


// Thread proc for the log thread
DWORD WINAPI Log_ThreadProc(LPVOID lpvGlobals)
{
	Log *pRemoteLog = static_cast<Log*>(lpvGlobals);

	while( !pRemoteLog->IsQuitting() )
	{
		pRemoteLog->DoAsynchronousUpdate();
	}

	return 0;
}


DWORD WINAPI Log_ListenerThreadProc(LPVOID lpvGlobals)
{
	Log *pRemoteLog = static_cast<Log*>(lpvGlobals);

	while( !pRemoteLog->IsQuitting() )
	{
		pRemoteLog->WaitForConnection();
		Sleep(50);
	}

	return 0;
}


Log::Log()
	: m_hSocket(0)
	, m_uiIDCounter(0)
	, m_uiSleepInterval(50)
	, m_bInitialized(false)
	, m_pAddrInfo(NULL)
	, m_pListeningAddrInfo(NULL)
	, m_bHasRemoteConnection(false)
	, m_hThread(0)
	, m_hListenerThread(0)
	, m_bQuitting(false)
{
}


void Log::Initialize()
{
	if( m_bInitialized )
		return;

	// We need to call ResetStream in order to write out placeholder length value
	ResetStream();

	// Create the update thread
	m_hThread = CreateThread(0, 0, Log_ThreadProc, static_cast<LPVOID>(this), CREATE_SUSPENDED, 0);
	m_hListenerThread = CreateThread(0, 0, Log_ListenerThreadProc, static_cast<LPVOID>(this), CREATE_SUSPENDED, 0);	

	WSADATA wsaData;
	// Initialize WinSock
	if (WSAStartup(MAKEWORD(2,0), &wsaData) != 0)
	{
		std::cout << "WSAStartup failed" << std::endl;
		exit(1);
	}

	struct addrinfo xHints;

	// First, load up address structs with getaddrinfo()
	memset(&xHints, 0, sizeof xHints);
	xHints.ai_family = AF_INET;
	xHints.ai_protocol = IPPROTO_TCP;
	xHints.ai_socktype = SOCK_STREAM;

	int iResult = getaddrinfo("127.0.0.1", "1337", &xHints, &m_pAddrInfo);
	iResult = getaddrinfo("127.0.0.1", "1338", &xHints, &m_pListeningAddrInfo);
	if( iResult != 0 )
	{
		return;
	}

	// Create two sockets: one for sending, and one for listening for connections
	m_hSocket = socket(m_pAddrInfo->ai_family, m_pAddrInfo->ai_socktype, m_pAddrInfo->ai_protocol);
	m_hListeningSocket = socket(m_pListeningAddrInfo->ai_family, m_pListeningAddrInfo->ai_socktype, m_pListeningAddrInfo->ai_protocol);

	// Set the socket to be non-blocking
	u_long ulNonBlocking = 1;
	iResult = ioctlsocket(m_hSocket, FIONBIO, &ulNonBlocking);

	if (iResult != NO_ERROR)
	{
		std::cout << "ioctlsocket failed with error: " << iResult;
	}

	m_bInitialized = true;

	// Connect
	ConnectToRemoteLogger();

	// Listen
	if( !m_bHasRemoteConnection )
	{
		iResult = bind(m_hListeningSocket, m_pListeningAddrInfo->ai_addr, m_pListeningAddrInfo->ai_addrlen);
		if( iResult != SOCKET_ERROR )
		{
		}
		else
		{
			int iLastError = WSAGetLastError();
			std::cout << iLastError;
		}

		iResult = listen(m_hListeningSocket, 128);
	}


	ResumeThread(m_hThread);
	ResumeThread(m_hListenerThread);
}


void Log::WaitForConnection()
{
	struct sockaddr_storage their_addr;
	socklen_t addr_size = sizeof their_addr;
	SOCKET resSock = accept(m_hListeningSocket, (struct sockaddr *)&their_addr, &addr_size);
	if( resSock != -1 )
	{
		ConnectToRemoteLogger();
	}
}


void Log::DoAsynchronousUpdate()
{
	if( m_bInitialized && m_bHasRemoteConnection )
	{
		SendData();
	}

	// Sleep to allow other threads to progress
	Sleep(m_uiSleepInterval);
}


void Log::ResetStream()
{
	// Clear the memory stream
	m_xMemoryStream.Clear();
	// Write a placeholder value, which will eventually get overwritten with the length of the stream
	m_xMemoryStream << (int)0;
}


void Log::SendData()
{
	m_csData.Enter();
	if( m_xMemoryStream.GetUsedSize() > sizeof(int) )
	{
		int iOffset = 0;

		// Grab the int* at the front of the memory stream
		int *pDataSize = (int*)m_xMemoryStream.GetData();
		// Set it to be the used size of the memory stream, minus its own size
		*pDataSize = m_xMemoryStream.GetUsedSize() - sizeof(int);

		// Send the memory stream
		int iBytesSent = send(m_hSocket, m_xMemoryStream.GetData(), m_xMemoryStream.GetUsedSize(), 0);

		// Now that we've sent all the data in our stream, reset it
		ResetStream();
	}
	m_csData.Leave();
}


void Log::ConnectToRemoteLogger()
{
	int iResult = 0;
		
	int iConnectionResult = connect(m_hSocket, m_pAddrInfo->ai_addr, m_pAddrInfo->ai_addrlen);
	m_bHasRemoteConnection = false;
	int iLastError = WSAGetLastError();

	fd_set setSocket;
	setSocket.fd_count = 1;
	setSocket.fd_array[0] = m_hSocket;

	struct timeval timeout;
	// Wait for 2 seconds for our connection attempt to succeed
	timeout.tv_sec = 2;
	timeout.tv_usec = 0;
	iResult = select( 0, NULL, &setSocket, NULL, &timeout );
	if( iResult < 0 )
	{
		iLastError = WSAGetLastError();
		std::cout << iLastError;
	}
	else
	if( iResult == 1 )
	{
		SendProcessInfo();

		m_bHasRemoteConnection = true;
	}
}


void Log::SendProcessInfo()
{
	// Send our process handle to the logging app
	m_csData.Enter();

	HANDLE hProcess = (HANDLE)GetCurrentProcessId();
	LogProcessInfo_s msgProcess;
	msgProcess.hProcess = (unsigned int)hProcess;

	TCHAR infoBuf[MAX_COMPUTERNAME_LENGTH+1];
	DWORD dwSize = MAX_COMPUTERNAME_LENGTH+1;
	GetComputerName(infoBuf, &dwSize);
	msgProcess.strComputerName = StringUtils::FromUnicode(infoBuf);

	msgProcess.Serialize(m_xMemoryStream);

	m_csData.Leave();
}


Log::~Log()
{
	m_bQuitting = true;

	// Close the sockets
	closesocket(m_hSocket);
	closesocket(m_hListeningSocket);
}


unsigned int Log::CreateID(const char *szTag, const char *szMessage, const unsigned char ucMsgClass)
{
	m_csInterface.Enter();

	++m_uiIDCounter;

	LogCreateIDMsg_s xMsg;

	xMsg.strTag = szTag;
	xMsg.strText = szMessage;
	xMsg.uiID = m_uiIDCounter;
	xMsg.ucMsgClass = ucMsgClass;

	// Store the format string
	std::string strFormatString = "";
	{
		for( size_t i = 0; i < strlen(szMessage); i++ )
		{
			if( szMessage[i] == '%' )
			{
				strFormatString = strFormatString + szMessage[i+1];
			}
		}
		m_mapTemplateInfo[xMsg.uiID].strParameterFormat = strFormatString;
	}

	m_mapTemplateInfo[xMsg.uiID].ucMsgClass = ucMsgClass;


	m_csData.Enter();
	xMsg.Serialize(m_xMemoryStream);
	m_csData.Leave();

	m_csInterface.Leave();

	return xMsg.uiID;
}


void Log::Add(unsigned int uiID, ...)
{
	m_csInterface.Enter();

	static USHORT (WINAPI *s_pfnCaptureStackBackTrace)(ULONG, ULONG, PVOID*, PULONG) = 0;

	if (s_pfnCaptureStackBackTrace == 0)
	{
		const HMODULE hNtDll = ::GetModuleHandle(L"ntdll.dll");
		reinterpret_cast<void*&>(s_pfnCaptureStackBackTrace) =
			::GetProcAddress(hNtDll, "RtlCaptureStackBackTrace");
	}

	int iNumEntries = 0;
	void *callStack[128];
	
	if( m_mapClassCallstackEnabled[m_mapTemplateInfo[uiID].ucMsgClass] )
	{
		iNumEntries = s_pfnCaptureStackBackTrace(0, 128, callStack, 0);
	}

	m_csData.Enter();
	
	LogNetMsg_s xMsg;
	xMsg.uiID = uiID;

	xMsg.Serialize(m_xMemoryStream);


	std::string strFormatString = m_mapTemplateInfo[uiID].strParameterFormat;


	if( !strFormatString.empty() )
	{
		va_list args;
		va_start(args, uiID);

		for( size_t i = 0; i < strFormatString.size(); i++ )
		{
			switch( strFormatString[i] )
			{
			case 'd':
				{
					int iVal = va_arg( args, int );
					m_xMemoryStream << iVal;
				}
				break;
			case 'f':
				{
					float fVal = (float)va_arg( args, double );
					m_xMemoryStream << fVal;
				}
				break;
			case 's':
				{
					char *szValue = va_arg( args, char * );
					m_xMemoryStream << szValue;
				}
				break;
			default:
				{
					// By default, treat it as an int
					int iVal = va_arg( args, int );
					m_xMemoryStream << iVal;
				}
				break;
			}
		}

		va_end(args);
	}

	m_xMemoryStream << iNumEntries;
	for( int i = 0; i < iNumEntries; i++ )
	{
		m_xMemoryStream << (unsigned int)callStack[i];
	}

	// Don't allow m_xMemoryStream to buffer too much data...
	int iMaxSize = 5000000;
	if( m_xMemoryStream.GetUsedSize() > iMaxSize )
	{
		ResetStream();
	}

	m_csData.Leave();
	m_csInterface.Leave();
}


void Log::CreateMessageClass(unsigned char ucID, const char *szName, const char *szTextColor, const char *szBackgroundColor, bool bBold, bool bCallstackEnabled)
{
	if( !m_bInitialized )
		return;

	LogCreateMsgClass_s xMsg;

	xMsg.strName = szName;
	xMsg.strTextColor = szTextColor;
	xMsg.strBackgroundColor = szBackgroundColor;
	xMsg.ucID = ucID;
	xMsg.bBold = bBold;

	m_csData.Enter();
	xMsg.Serialize(m_xMemoryStream);
	m_csData.Leave();

	m_mapClassCallstackEnabled[ucID] = bCallstackEnabled;
}


void LogNetMsg_s::Serialize(MemoryStream &xMemoryStream)
{
	xMemoryStream << ucMsgType;
	xMemoryStream << uiID;
}


void LogNetMsg_s::Deserialize(const MemoryStream &xMemoryStream)
{
	xMemoryStream >> ucMsgType;
	xMemoryStream >> uiID;
}


void LogCreateIDMsg_s::Serialize(MemoryStream &xMemoryStream)
{
	xMemoryStream << ucMsgType;
	xMemoryStream << uiID;
	xMemoryStream << ucMsgClass;
	xMemoryStream << strText;
	xMemoryStream << strTag;
}


void LogCreateIDMsg_s::Deserialize(const MemoryStream &xMemoryStream)
{
	xMemoryStream >> ucMsgType;
	xMemoryStream >> uiID;
	xMemoryStream >> ucMsgClass;
	xMemoryStream >> strText;
	xMemoryStream >> strTag;
}


void LogCreateMsgClass_s::Serialize(MemoryStream &xMemoryStream)
{
	xMemoryStream << ucMsgType;
	xMemoryStream << ucID;
	xMemoryStream << strName;
	xMemoryStream << strTextColor;
	xMemoryStream << strBackgroundColor;
	xMemoryStream << bBold;
}


void LogCreateMsgClass_s::Deserialize(const MemoryStream &xMemoryStream)
{
	xMemoryStream >> ucMsgType;
	xMemoryStream >> ucID;
	xMemoryStream >> strName;
	xMemoryStream >> strTextColor;
	xMemoryStream >> strBackgroundColor;
	xMemoryStream >> bBold;
}


void LogProcessInfo_s::Serialize(MemoryStream &xMemoryStream)
{
	xMemoryStream << ucMsgType;
	xMemoryStream << hProcess;
	xMemoryStream << strComputerName;
}


void LogProcessInfo_s::Deserialize(const MemoryStream &xMemoryStream)
{
	xMemoryStream >> ucMsgType;
	xMemoryStream >> hProcess;
	xMemoryStream >> strComputerName;
}